#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
更新 samples.json 中的 sample_type（dataset_type）字段

功能：
1. 从 test.csv 中读取所有测试集的 scenario_id
2. 将 samples.json 中对应的 scenario 的 dataset_type 设置为 "test"
3. 其余 scenario 的 dataset_type 设置为 "train"
4. 更新 metadata 中的统计信息

使用方法：
    python scripts/update_sample_types.py --test-csv data/envs/aiops-22/test.csv --samples-json data/envs/aiops-22/samples.json
"""

import argparse
import csv
import json
import sys
from pathlib import Path
from typing import Set


def load_test_scenario_ids(test_csv_path: str) -> Set[str]:
    """
    从 test.csv 文件中加载所有测试集的 scenario_id
    
    参数:
        test_csv_path: test.csv 文件路径
        
    返回:
        测试集 scenario_id 的集合
    """
    test_scenario_ids = set()
    
    with open(test_csv_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            scenario_id = row.get('scenario_id', '').strip()
            if scenario_id:  # 跳过空行
                test_scenario_ids.add(scenario_id)
    
    return test_scenario_ids


def update_samples_json(samples_json_path: str, test_scenario_ids: Set[str]) -> None:
    """
    更新 samples.json 文件中的 dataset_type 和 sample_type 字段

    参数:
        samples_json_path: samples.json 文件路径
        test_scenario_ids: 测试集 scenario_id 的集合
    """
    # 读取 samples.json
    with open(samples_json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 统计信息
    train_count = 0
    test_count = 0
    train_samples = 0
    test_samples = 0

    # 第一步：先将所有 scenario 的 dataset_type 设置为空
    scenarios = data.get('scenarios', [])
    for scenario in scenarios:
        scenario['dataset_type'] = ''

    # 第二步：先将所有 sample 的 sample_type 设置为空
    samples = data.get('samples', [])
    for sample in samples:
        sample['sample_type'] = ''

    print("已将所有场景的 dataset_type 和样本的 sample_type 重置为空")

    # 第三步：根据 test.csv 将场景和样本的 dataset_type/sample_type 设置为 train/test
    for scenario in scenarios:
        scenario_id = scenario.get('scenario_id', '')
        trace_count = scenario.get('trace_count', 0)

        if scenario_id in test_scenario_ids:
            # 如果在测试集中，设置为 test
            scenario['dataset_type'] = 'test'
            test_count += 1
            test_samples += trace_count

            # 同时将该场景下的所有样本设置为 test
            for sample in samples:
                if sample.get('scenario_id') == scenario_id:
                    sample['sample_type'] = 'test'

            print(f"设置测试集场景: {scenario_id} (样本数: {trace_count})")
        else:
            # 否则设置为 train
            scenario['dataset_type'] = 'train'
            train_count += 1
            train_samples += trace_count

            # 同时将该场景下的所有样本设置为 train
            for sample in samples:
                if sample.get('scenario_id') == scenario_id:
                    sample['sample_type'] = 'train'

            print(f"设置训练集场景: {scenario_id} (样本数: {trace_count})")

    # 更新 metadata 统计信息
    if 'metadata' in data:
        data['metadata']['train_scenarios'] = train_count
        data['metadata']['test_scenarios'] = test_count
        data['metadata']['train_samples'] = train_samples
        data['metadata']['test_samples'] = test_samples
        data['metadata']['total_scenarios'] = train_count + test_count
        data['metadata']['total_samples'] = train_samples + test_samples

    # 保存更新后的 samples.json
    with open(samples_json_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    print(f"\n更新完成！")
    print(f"训练集场景数: {train_count}, 样本数: {train_samples}")
    print(f"测试集场景数: {test_count}, 样本数: {test_samples}")
    print(f"总计场景数: {train_count + test_count}, 样本数: {train_samples + test_samples}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='根据 test.csv 更新 samples.json 中的 dataset_type 字段'
    )
    parser.add_argument(
        '--test-csv',
        type=str,
        default='data/envs/aiops-22/test.csv',
        help='test.csv 文件路径（默认: data/envs/aiops-22/test.csv）'
    )
    parser.add_argument(
        '--samples-json',
        type=str,
        default='data/envs/aiops-22/samples.json',
        help='samples.json 文件路径（默认: data/envs/aiops-22/samples.json）'
    )
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    test_csv_path = Path(args.test_csv)
    samples_json_path = Path(args.samples_json)
    
    if not test_csv_path.exists():
        print(f"错误: 文件不存在: {test_csv_path}", file=sys.stderr)
        sys.exit(1)
    
    if not samples_json_path.exists():
        print(f"错误: 文件不存在: {samples_json_path}", file=sys.stderr)
        sys.exit(1)
    
    # 加载测试集 scenario_id
    print(f"正在从 {test_csv_path} 加载测试集 scenario_id...")
    test_scenario_ids = load_test_scenario_ids(str(test_csv_path))
    print(f"找到 {len(test_scenario_ids)} 个测试集 scenario_id")
    
    # 更新 samples.json
    print(f"\n正在更新 {samples_json_path}...")
    update_samples_json(str(samples_json_path), test_scenario_ids)


if __name__ == '__main__':
    main()

